# %%
from train import *
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "6"



# %%
import joblib
f= open('./Data/Processed/customer_feat.pkl','rb')
customer_id = joblib.load(f)
FN = joblib.load(f)
Active = joblib.load(f)
club_member_status = joblib.load(f)
fashion_news_frequency = joblib.load(f)
age = joblib.load(f)
customer_cnt = joblib.load(f)
customer_bidict = joblib.load(f)

# %%
model, feature_columns, behavior_feature_list = set_model()

# %%
model.load_state_dict(torch.load("./Model/model2.pt"))

# %%
pos_artid_seq_last = joblib.load('./Data/Processed/divide/pos_artid_seq_val.db')
pos_pcode_seq_last = joblib.load('./Data/Processed/divide/pos_pcode_seq_val.db')

# %%
pos_feat = open('./Data/Processed/divide/pos_feat_val.db','rb')
pos_custom_id_last = joblib.load(pos_feat)
pos_FN_last = joblib.load(pos_feat)
pos_Active_last = joblib.load(pos_feat)
pos_club_member_status_last = joblib.load(pos_feat)
pos_fashion_news_frequency_last = joblib.load(pos_feat)
pos_age_last = joblib.load(pos_feat)
pos_artid_last = joblib.load(pos_feat)
pos_pcode_last = joblib.load(pos_feat)

# %%
from bidict import bidict

# %%
pos_in_last = np.array([-1 for i in range(1371980)])
for i in range(pos_custom_id_last.size):
    pos_in_last[pos_custom_id_last[i]] = i


# %%
history = {}
for i in customer_id:
    if pos_in_last[i] != -1:
        artid_seq_i = pos_artid_seq_last[pos_in_last[i]]
        pcode_seq_i = pos_pcode_seq_last[pos_in_last[i]]
        if artid_seq_i[-1] == 0:
            for j in range(200):
                if artid_seq_i[j] == 0:
                    artid_seq_i[j] = pos_artid_last[pos_in_last[i]]
                    break
        else:
            artid_seq_i = np.append(artid_seq_i[1:],pos_artid_last[pos_in_last[i]])
            pcode_seq_i = np.append(pcode_seq_i[1:],pos_pcode_last[pos_in_last[i]])
        history[i] = {'artid_seq':artid_seq_i,'pcode_seq':pcode_seq_i}
    else:
        artid_seq_i = np.array([0 for j in range(200)])
        pcode_seq_i = np.array([0 for j in range(200)])
        history[i] = {'artid_seq':artid_seq_i,'pcode_seq':pcode_seq_i}

# %%
article_sparse = open('./Data/Processed/article_sparse.pkl','rb')
artid = joblib.load(article_sparse)
pcode = joblib.load(article_sparse)
artid_bidict = joblib.load(article_sparse)
pcode_bidict = joblib.load(article_sparse)
artid_cnt = joblib.load(article_sparse)

# %%
article_pool = []
topk_index_100 = open('./Data/Processed/topk_index_100.pkl','rb')
topk_index = joblib.load(topk_index_100)
defaultsample = np.random.choice(range(1, 105543), 100)
print(customer_id)
for i in customer_id:
    if (i % 10000 == 0):
        print(f'{i} w ok')
    if history[i]['artid_seq'][0] == 0:
        history[i]['artid_seq'][0] = np.random.randint(1,105543)
    t = 0
    for j in range(200):
        if history[i]['artid_seq'][j] != 0 and j!=199:
            t = history[i]['artid_seq'][j]
        elif history[i]['artid_seq'][j] == 0:
            article_pool.append(topk_index[t - 1] + 1)
            break 
        elif j == 199:
            article_pool.append(topk_index[history[i]['artid_seq'][j] - 1] + 1)
print(len(article_pool))

# %%


# %%



# %%
#result = []
#batch = {name:np.array([]) for name in get_feature_names(feature_columns)}
#article_pool[i] = article_pool[i].astype(int)
#num = article_pool[i].size
import csv
Output = open('result4.csv','w')
writer = csv.writer(Output)
header = ['customer_id','article0','article1','article2','article3','article4','article5','article6','article7','article8','article9','article10','article11']

for i in customer_id:
    num = 100
    batch = {}
    if(i % 100000 == 0):
        print(i)
    batch['pos_custom_id'] = np.array([i]).repeat(num)
    batch['pos_Active'] = np.array([Active[i]]).repeat(num)
    batch['pos_FN'] = np.array([FN[i]]).repeat(num)
    batch['pos_club_member_status'] = np.array([club_member_status[i]]).repeat(num)
    batch['pos_fashion_news_frequency'] = np.array([fashion_news_frequency[i]]).repeat(num)
    batch['pos_artid'] = article_pool[i]
    batch['pos_pcode'] = pcode[article_pool[i] - 1]
    batch['pos_sales_channel_id'] = np.random.randint(2,size = num)
    batch['pos_age'] = np.array([age[i] for j in range(num)])
    batch['hist_pos_artid'] = np.array([history[i]['artid_seq'] for j in range(num)])
    #seq_length = np.count_nonzero(history[i]['artid_seq'])
    batch['seq_length'] = np.array([np.count_nonzero(history[i]['artid_seq'])]).repeat(num)
    batch['hist_pos_pcode'] = np.array([history[i]['pcode_seq'] for j in range(num)])
#Batch.append(batch)
# if(i % 25 == 0):
#     ctr = model.predict()
    ctr = model.predict(batch, num)  
    ctr = ctr.squeeze(1)
    #print(ctr) 
    #print(np.argsort(ctr)[-12:])
    result = article_pool[i][np.argsort(ctr)[-12:]]
    #print(result)
    writer.writerow([i] + result.tolist())     
    #result.append(article_pool[i][np.argsort(ctr)[-12:]])


# %%